package com.yiyefanchen.ai.utils;

import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * SSE 工具类，用于管理 SseEmitter 连接和推送消息
 */
public class SseEmitterUtil {

    // 存储所有 SSE 连接的缓存（Key 可以是用户ID或会话ID）
    private static final Map<String, SseEmitter> sseCache = new ConcurrentHashMap<>();


    /**
     * 创建 SSE 连接
     * @param id 客户端唯一标识（如用户ID或会话ID）
     * @return SseEmitter
     */
    public static SseEmitter createSse(String id) {
        // 设置超时时间（0表示永不超时，按需调整）
        SseEmitter sseEmitter = new SseEmitter(0L);

        // 注册回调：连接关闭时移除缓存
        sseEmitter.onCompletion(() -> removeSse(id));
        sseEmitter.onError((e) -> removeSse(id));
        sseEmitter.onTimeout(() -> removeSse(id));

        // 存入缓存
        sseCache.put(id, sseEmitter);
        return sseEmitter;
    }

    /**
     * 向指定客户端推送消息
     * @param id 客户端标识
     * @param message 消息内容
     */
    public static void sendMessage(String id, Object message) {
        SseEmitter sseEmitter = sseCache.get(id);
        if (sseEmitter != null) {
            try {
                sseEmitter.send(SseEmitter.event().data(message));
            } catch (IOException e) {
                removeSse(id); // 推送失败时移除连接
                throw new RuntimeException("SSE 推送失败: " + e.getMessage());
            }
        }
    }

    /**
     * 移除 SSE 连接
     * @param id 客户端标识
     */
    public static void removeSse(String id) {
        SseEmitter sseEmitter = sseCache.get(id);
        if (sseEmitter != null) {
            sseEmitter.complete();
            sseCache.remove(id);
        }
    }

    /**
     * 获取当前所有活跃连接
     */
    public static Map<String, SseEmitter> getSseCache() {
        return sseCache;
    }
}